import cv2
from PIL import Image, ImageDraw, ImageFont
import sys
CENTERNET_PATH = '/home/predict'
sys.path.insert(0, CENTERNET_PATH)

from yolov4 import init_yolov4,detect,load_class_names,plot_boxes_cv2
import numpy as np
import math


from datetime import datetime

def cv2ImgAddText(img, text, left, top, textColor=(0, 255, 0), textSize=20):
    if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
        img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    # 创建一个可以在给定图像上绘图的对象
    draw = ImageDraw.Draw(img)
    # 字体的格式
    fontStyle = ImageFont.truetype(
        "SimHei.ttf", textSize, encoding="utf-8")
    # 绘制文本
    draw.text((left, top), text, textColor, font=fontStyle)
    # 转换回OpenCV格式
    return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)


def plot_boxes_cv2_zh(img, boxes, conf_thresh=0.4,savename=None, class_names=None, color=None):
    import cv2
    img = np.copy(img)
    colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)

    def get_color(c, x, max_val):
        ratio = float(x) / max_val * 5
        i = int(math.floor(ratio))
        j = int(math.ceil(ratio))
        ratio = ratio - i
        r = (1 - ratio) * colors[i][c] + ratio * colors[j][c]
        return int(r * 255)

    width = img.shape[1]
    height = img.shape[0]
    res=[]
    for i in range(len(boxes)):
        box = boxes[i]
        if len(box) >= 7 and box[5]>conf_thresh:      
            x1 = int((box[0] - box[2] / 2.0) * width)
            y1 = int((box[1] - box[3] / 2.0) * height)
            x2 = int((box[0] + box[2] / 2.0) * width)
            y2 = int((box[1] + box[3] / 2.0) * height)
            rgb_red = (0,0,255)  #红
            rgb_green = (0,255,0) #绿
            rgb_blue = (255,255,0)  #蓝
            rgb_black = (0,0,0)       #黑
            rgb_white = (255,255,255) #白
            rgb_yellow = (0,255,255) #黄

            if color:
                rgb = color
            else:
                rgb = (255, 0, 0)
            if len(box) >= 7 and class_names:
                cls_conf = box[5]
                cls_id = box[6]
                classes = len(class_names)
                offset = cls_id * 123457 % classes
                red = get_color(2, offset, classes)
                green = get_color(1, offset, classes)
                blue = get_color(0, offset, classes)
                if color is None:
                    rgb = (red, green, blue)
           # img = cv2.putText(img, class_names[cls_id], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1.2, rgb, 1)
                text=class_names[cls_id]
                if text=='fire':
                    rgb = rgb_red
                    text='火焰'
                elif text=='gang':
                    rgb = rgb_blue
                    text='钢'   
                elif text=='red':
                    rgb = rgb_red
                    text='红'   
                elif text=='yellow':
                    rgb = rgb_yellow
                    text='黄'   
                elif text=='blue':
                    rgb = rgb_blue
                    text='蓝'
                elif text=='white':
                    rgb = rgb_white
                    text='白'    
                elif text=='worker':
                    rgb = rgb_blue
                elif text=='person':
                    rgb = rgb_black
                elif text=='firehua':
                    rgb = rgb_red
                    text='火苗'   
                if text!='worker' and text!='person' and text!='火焰':
                    if y1>20:
                        img = cv2ImgAddText(img, text, x1, y1-20, (0, 0 , 139), 20) 
                    else :
                        img = cv2ImgAddText(img, text, x1, 0, (0, 0 , 139), 20) 
                if text=='火焰':  #火焰会打出置信度
                    text=text+' '+str(cls_conf)
                    img = cv2ImgAddText(img, text, x1, y1-20, (0, 0 , 139), 20) 
            img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, 1)
            res.append([(x1, y1), (x2, y2),cls_conf,class_names[cls_id]])#(x1, y1), (x2, y2),conf,name
    if savename:
        cv2.imwrite(savename, img)
    return img,res


def write_txt(id_list,cam_num,print_txt):#用于记录需要测试的日志
    with open('txt/'+str(id_list[cam_num-1])+'.txt','a+',encoding='utf-8') as f:
        f.write(str(datetime.now())+' :'+print_txt+'\n')


def electron_frame(frame,size_frame,electron_fence):   #绘制电子围栏和电子屏蔽
    if electron_fence!='None':
        electrons=electron_fence.split('-') #使用;分割
        for electron in electrons:
            electron=electron[1:-1].split(';')
            if electron[0]=='0':  #电子屏蔽
                frame = cv2.rectangle(frame, (int(float(electron[1])*size_frame[0]), int(float(electron[3])*size_frame[1])), (int(float(electron[2])*size_frame[0]), int(float(electron[4])*size_frame[1])), (0,255,0), 2)
            elif electron[0]=='1': #电子围栏
                frame = cv2.rectangle(frame, (int(float(electron[1])*size_frame[0]), int(float(electron[3])*size_frame[1])), (int(float(electron[2])*size_frame[0]), int(float(electron[4])*size_frame[1])), (0,0,255), 2)
    return frame










